package project.Components;

import jakarta.websocket.*;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.annotation.CrossOrigin;
import project.Config.HttpSessionConfig;
import project.Pojo.Chat.AI.ChatRecord;
import project.Service.ChatService;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint(value = "/chat/ai/{id}", configurator = HttpSessionConfig.class)
@CrossOrigin
public class AiWebSocketServer {

    //暂存用户名（通过地址传入)
    String username;

    //传入模型的提示词
    SystemMessage systemMessage;

    //从上下文获取bean
    static private ApplicationContext applicationContext;

    static public void setApplicationContext(ApplicationContext applicationContext) {
        AiWebSocketServer.applicationContext = applicationContext;
    }

    private ChatService chatService;
    private OllamaChatModel ollama;

    //存放当前在线连接的用户信息(session池)
    private static final ConcurrentHashMap<String, Session> sessionPool = new ConcurrentHashMap<>();

    //建立websocket连接时调用
    @OnOpen
    public void onOpen(Session session, EndpointConfig config, @PathParam("id") final String id) throws IOException {
        username = id;
        //暂存用户session,断开连接时在onClose删除
        sessionPool.put(username, session);

        //设置提示信息
        systemMessage = new SystemMessage("""
                你是一名专业的心理咨询师，擅长认知行为疗法（CBT）。
                请用温暖、共情的中文与用户交流，避免直接给出建议，
                而是通过提问引导用户自我反思。语气要温和且非评判性。
                """);

        ollama = applicationContext.getBean(OllamaChatModel.class);
        chatService = applicationContext.getBean(ChatService.class);
    }


    //收到浏览器发送的消息时自动调用
    @OnMessage
    public void onMessage(String message) throws IOException {
        
        //存聊天记录
        ChatRecord userRecord = new ChatRecord();
        userRecord.setMsgDate(new Date());
        userRecord.setMsgContent(message);
        userRecord.setUserId(username);
        //0表示用户，1表示ai
        userRecord.setMsgRole(0);
        chatService.insertAiRecord(userRecord);


        //用户输入
        UserMessage userMessage = new UserMessage(message);
        //创建提示,拼接提示词和用户信息
        Prompt prompt = new Prompt(List.of(systemMessage, userMessage));
        ChatResponse response = ollama.call(prompt);
        //获取session并发送AI响应
        String content = response.getResult().getOutput().getText();


        //存聊天记录
        ChatRecord ai = new ChatRecord();
        ai.setMsgDate(new Date());
        ai.setMsgContent(content);
        ai.setUserId(username);
        //0表示用户，1表示ai
        ai.setMsgRole(1);
        chatService.insertAiRecord(ai);


        sessionPool.get(username).getBasicRemote().sendText(content);
    }


    @OnClose
    public void onClose(Session session) throws IOException {

        sessionPool.remove(username);

    }
}
